#ERA5数据预处理（垂直方向和水平方向插值）并插值到FY-4A的网格中 -*-
import xarray as xr
import numpy as np
import pandas as pd
import multiprocessing as mp
import os
import traceback
import sys

# ==============================================================================
# 1. 配置 (Configuration)
# ==============================================================================
ERA5_SURFACE_PATH = '/mnt/datastore/liudddata/ERA5/temp202005_06.nc'
ERA5_Q_PATH = '/mnt/datastore/liudddata/ERA5/validation_plots/spacific_humidity_20200601_24.nc'
COORD_FILE_PATH = '/home/liudd/data_preprocessing/FY4A_coordinates.nc'
OUTPUT_DIR = '/mnt/datastore/liudddata/ERA5_output/2020005_06_all_final/'
FAILED_LOG_PATH = os.path.join(OUTPUT_DIR, "failed_timestamps.txt")

# 创建输出目录
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 读取目标网格
try:
    with xr.open_dataset(COORD_FILE_PATH) as coord_ds:
        lon_fy4a = coord_ds['lon'].values.T
        lat_fy4a = coord_ds['lat'].values.T
except Exception as e:
    print(f"错误：读取FY4A坐标文件失败: {e}")
    sys.exit()

# 预计算目标坐标
lon_fy4a_360 = lon_fy4a % 360
target_lon = xr.DataArray(lon_fy4a_360, dims=('y', 'x'))
target_lat = xr.DataArray(lat_fy4a, dims=('y', 'x'))
print("--- 坐标预计算完成 ---")


# ==============================================================================
# 2. 自定义垂直插值/边界赋值函数 (Custom Vertical Interpolation)
# ==============================================================================
def interpolate_or_clamp(q_profile, log_level_profile, log_sp_scalar):
    """
    对单个垂直剖面进行插值或边界赋值。
    如果 log_sp 超出 log_levels 范围，则使用最近的边界值。
    """
    # 检查输入有效性
    if np.isnan(log_sp_scalar) or len(log_level_profile) < 1 or np.all(np.isnan(q_profile)):
        return np.nan

    # 使用边界值
    if log_sp_scalar <= log_level_profile[0]:
        return q_profile[0]
    elif log_sp_scalar >= log_level_profile[-1]:
        return q_profile[-1]
    # 在范围内则进行线性插值
    else:
        return np.interp(log_sp_scalar, log_level_profile, q_profile)


# ==============================================================================
# 3. 核心处理函数 (Core Processing Function)
# ==============================================================================
def process_time_step(timestamp):
    """
    处理单个时间步的ERA5数据。
    """
    time_str = pd.to_datetime(timestamp).strftime('%Y-%m-%d %H:%M')
    print(f"\n--- 正在处理: {time_str} ---")

    try:
        # 加载数据
        with xr.open_dataset(ERA5_SURFACE_PATH) as surface_ds, xr.open_dataset(ERA5_Q_PATH) as q_ds:
            sp_da = surface_ds['sp'].sel(valid_time=timestamp)
            t2m_da = surface_ds['t2m'].sel(valid_time=timestamp)

            if 'pressure_level' in q_ds.dims:
                q_ds = q_ds.rename({'pressure_level': 'level'})
            q_multilayer_da = q_ds['q'].sel(valid_time=timestamp)

        if q_multilayer_da.level.size < 2 or np.isnan(q_multilayer_da.values).all():
            raise ValueError(f"比湿层数不足或数据无效（层数: {q_multilayer_da.level.size}）")

        # 准备垂直插值所需数据
        sp_hpa = sp_da / 100.0
        log_sp = np.log(sp_hpa)

        if np.any(np.isnan(log_sp.values)) or np.any(np.isinf(log_sp.values)):
            raise ValueError("计算 log(sp) 时出现 NaN 或 Inf")

        # 确保 level 维度是第一个，以便 apply_ufunc 正确处理
        q_data = q_multilayer_da.transpose("level", "latitude", "longitude")
        log_levels = np.log(q_data.level.values)

        # 使用 apply_ufunc 进行逐点垂直插值
        # 这是你成功的核心实现，非常棒！
        q_surface = xr.apply_ufunc(
            interpolate_or_clamp,
            q_data,
            xr.DataArray(log_levels, dims=['level']), # 将 log_levels 也作为 DataArray 传入
            log_sp,
            input_core_dims=[["level"], ["level"], []], # 确保维度匹配
            output_core_dims=[[]],
            exclude_dims=set(('level',)), # 排除 level 维度，因为它在计算后被移除
            vectorize=True,
            dask="parallelized",
            output_dtypes=[q_data.dtype]
        )

        if np.isnan(q_surface.values).all():
            raise ValueError("垂直插值后的比湿全为 NaN")

        # 水平插值至 FY4A 格点
        t2m_interp_da = t2m_da.interp(longitude=target_lon, latitude=target_lat, method="linear")
        sp_interp_da = sp_da.interp(longitude=target_lon, latitude=target_lat, method="linear")
        q_surface_interp_da = q_surface.interp(longitude=target_lon, latitude=target_lat, method="linear")

        # ==================================================================
        #  优化点 (OPTIMIZATION)
        # ==================================================================
        # 将秒精度的时间戳转换为纳秒精度，以消除 xarray 的 UserWarning
        timestamp_ns = pd.to_datetime(timestamp).to_datetime64()

        # 构建输出
        ds_out = xr.Dataset(
            {
                "temp_2m": (("y", "x"), t2m_interp_da.values),
                "surface_pressure": (("y", "x"), sp_interp_da.values),
                "surface_specific_humidity": (("y", "x"), q_surface_interp_da.values)
            },
            coords={"time": timestamp_ns, "lon": (("y", "x"), lon_fy4a), "lat": (("y", "x"), lat_fy4a)}
        )

        output_filename = os.path.join(
            OUTPUT_DIR,
            f'ERA5_{pd.to_datetime(timestamp).strftime("%Y%m%d%H")}_surface_vars_FY4A_grid.nc'
        )
        ds_out.to_netcdf(output_filename)
        print(f"--- ✅ 成功保存: {os.path.basename(output_filename)} ---")

    except Exception as e:
        print(f"!!!!!! 时间 {time_str} 处理失败: {e} !!!!!!")
        with open(FAILED_LOG_PATH, 'a') as f:
            f.write(f"{time_str} | {str(e)}\n")
        traceback.print_exc(file=sys.stdout)


# ==============================================================================
# 4. 主程序入口 (Main Execution)
# ==============================================================================
if __name__ == '__main__':
    print("--- 正在进行时间戳预匹配... ---")
    try:
        with xr.open_dataset(ERA5_SURFACE_PATH) as ds_surf, xr.open_dataset(ERA5_Q_PATH) as ds_q:
            surface_times = ds_surf.valid_time.values
            q_times = ds_q.valid_time.values
        valid_timestamps = np.intersect1d(surface_times, q_times)
        if len(valid_timestamps) == 0:
            print("!!!!!! 严重错误: 没有共同的时间戳，程序退出。 !!!!!!")
            sys.exit()
        print(f"--- 匹配完成，共 {len(valid_timestamps)} 个时间点 ---")
    except Exception as e:
        print(f"!!!!!! 时间戳预处理失败: {e} !!!!!!")
        sys.exit()

    # 建议使用稍小一些的进程数，以避免过多I/O竞争
    num_processes = min(mp.cpu_count(), 8)
    print(f"--- 启动 {num_processes} 个进程并行处理 ---")
    with mp.Pool(processes=num_processes) as pool:
        pool.map(process_time_step, valid_timestamps)

    print("\n--- ✅ 所有处理完成 ---")
